# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
r"""Downloads and converts Flowers data to TFRecords of TF-Example protos.

This module downloads the Flowers data, uncompresses it, reads the files
that make up the Flowers data and creates two TFRecord datasets: one for train
and one for test. Each TFRecord dataset is comprised of a set of TF-Example
protocol buffers, each of which contain a single image and label.

The script should take about a minute to run.

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from multiprocessing import cpu_count

import math
import os
import random
import sys
import time
import pickle
import shutil
import threadpool
import _thread
import tensorflow as tf

sys.path.append("./")
import dataset_utils

OUTPUT_PATH = "./output_tfrecord"
PREFIX = ""  # should be train or validation
_RANDOM_SEED = 0


class TfMd5:
    def __init__(self, dump_file_name):
        self.tf_md5_list = []
        self.dump_file_name = dump_file_name

    def add_new_record(self, filename):
        self.tf_md5_list.append(os.popen("md5sum " + filename).read().split(' ')[0])

    def dump(self):
        f = open(self.dump_file_name, 'wb')
        pickle.dump(self.tf_md5_list, f)
        f.close()

    def load(self):
        f = open(self.dump_file_name, 'rb')
        self.tf_md5_list = pickle.load(f)


class GlobalProcess:
    def __init__(self):
        self.all_step = 0
        self.finished_file = []

    def set_all_step(self, step):
        self.all_step = step

    def add_finish(self, file):
        self.finished_file.append(file)


class ImageReader(object):
    """Helper class that provides TensorFlow image coding utilities."""

    def __init__(self):
        # Initializes function that decodes RGB JPEG data.
        self._decode_jpeg_data = tf.placeholder(dtype=tf.string)
        self._decode_jpeg = tf.image.decode_jpeg(self._decode_jpeg_data, channels=3)

    def read_image_dims(self, sess, image_data):
        image = self.decode_jpeg(sess, image_data)
        return image.shape[0], image.shape[1]

    def decode_jpeg(self, sess, image_data):
        image = sess.run(self._decode_jpeg,
                         feed_dict={self._decode_jpeg_data: image_data})
        assert len(image.shape) == 3
        assert image.shape[2] == 3
        return image


def check_label_same(label_file, labels_to_class_names):
    if os.path.exists(label_file) is False:
        return False
    read_class_names = []
    f = open(label_file)
    line = " "
    while line:
        line = f.readline().strip()
        if len(line.split(':')) >= 2:
            read_class_names.append(line.split(':')[1])
    f.close()
    file_labels_to_class_names = dict(zip(range(len(read_class_names)), read_class_names))
    if file_labels_to_class_names == labels_to_class_names:
        return True
    else:
        return False


def get_pkl_list(dir):
    pkl_file_list = []
    for filename in os.listdir(dir):
        if os.path.splitext(filename)[1] == ".pkl":
            pkl_file_list.append(dir + "/" + filename)
    return pkl_file_list


def check_pair_of_pkl_tf(outpath):
    pkl_file_list = []
    for filename in os.listdir(outpath):
        path = os.path.join(outpath, filename)
        if os.path.splitext(path)[1] == ".tfrecord":
            if os.path.exists(os.path.splitext(path)[0] + ".pkl") is False:
                os.remove(path)
        if os.path.splitext(path)[1] == ".pkl":
            if os.path.exists(os.path.splitext(path)[0] + ".tfrecord") is False:
                os.remove(path)
            else:
                pkl_file_list.append(path)
    return pkl_file_list


def islistA_contain_in_listB(listA, listB):
    for A in listA:
        if A in listB:
            continue
        else:
            return False
    return True


def fill_dict_and_list(p_file,photo_file_md5_dict,md5_duplicated_file_list,photo_file_md5_list):
    p_file_md5 = os.popen("md5sum " + p_file).read().split(' ')[0]
    if photo_file_md5_dict.get(p_file_md5) != None:
        md5_duplicated_file_list.append(p_file)
    else:
        photo_file_md5_dict[p_file_md5] = p_file
        photo_file_md5_list.append(p_file_md5)


def wipe_out_already_finished_task(photo_files, pkl_file_list):
    if len(pkl_file_list) == 0:
        return photo_files
    photo_file_md5_list = []
    photo_file_md5_dict = {}
    any_pkl_removed = False
    md5_duplicated_file_list = []
    cnt = 0
    for p_file in photo_files:
        #fill_dict_and_list(p_file,photo_file_md5_dict,md5_duplicated_file_list,photo_file_md5_list)
        _thread.start_new_thread(fill_dict_and_list,(p_file,photo_file_md5_dict,md5_duplicated_file_list,photo_file_md5_list))
        sys.stdout.write("\r >> Fill dict list thread: %d " % cnt)
        cnt +=1
    print("")
    fill_dict_and_list(photo_files[-1], photo_file_md5_dict, md5_duplicated_file_list, photo_file_md5_list)
    print(' CHECKING build dict AT: %.5f second' % (time.time() - start_time))
    for file in md5_duplicated_file_list:
        photo_files.remove(file)

    for pkl_file in pkl_file_list:
        md5_con = TfMd5(pkl_file)
        md5_con.load()
        if islistA_contain_in_listB(md5_con.tf_md5_list, photo_file_md5_list) is False:
            any_pkl_removed = True
            os.remove(pkl_file)
            continue
        else:
            for md5_key in md5_con.tf_md5_list:
                if photo_file_md5_dict.get(md5_key) != None:
                    if photo_files.count(photo_file_md5_dict[md5_key]) > 0:
                        photo_files.remove(photo_file_md5_dict[md5_key])
    print(' CHECKING restore process AT: %.5f second' % (time.time() - start_time))
    if len(photo_files) != 0 and any_pkl_removed is False:
        md5_con = TfMd5(pkl_file_list[-1])
        md5_con.load()
        for md5 in md5_con.tf_md5_list:
            if photo_file_md5_dict.get(md5) !=None:
                append_file = photo_file_md5_dict[md5]
                if photo_files.count(append_file) == 0:
                    photo_files.append(append_file)
                else:
                    print(" photo file already has this file")
            else:
                print(" md5 not in photo file md5 dict")
        os.remove(pkl_file_list[-1])
    print(' CHECKING only adding file AT: %.5f second' % (time.time() - start_time))
    return photo_files


def _get_filenames_and_classes(dataset_dir):
    """Returns a list of filenames and inferred class names.

    Args:
      dataset_dir: A directory containing a set of subdirectories representing
        class names. Each subdirectory should contain PNG or JPG encoded images.

    Returns:
      A list of image file paths, relative to `dataset_dir` and the list of
      subdirectories, representing class names.
    """
    file_root = os.path.join(dataset_dir, './')
    directories = []
    class_names = []
    for filename in os.listdir(file_root):
        path = os.path.join(file_root, filename)
        if os.path.isdir(path):
            directories.append(path)
            class_names.append(filename)

    photo_filenames = []
    for directory in directories:
        for filename in os.listdir(directory):
            path = os.path.join(directory, filename)
            photo_filenames.append(path)

    return photo_filenames, sorted(class_names)


def _get_dataset_filename(dir, split_name, shard_id, num_shards):
    output_filename = '%s_%05d_of_%05d_%d.tfrecord' % (
        split_name, shard_id, num_shards, VERSION)
    return os.path.join(dir, output_filename)


def create_tf(split_name, filenames, class_names_to_ids, dataset_dir,
              num_shards, shard_id, num_per_shard, image_reader, sess):
    output_filename = _get_dataset_filename(
        OUTPUT_PATH, split_name, shard_id, num_shards)
    md5_seri = TfMd5(os.path.splitext(output_filename)[0] + ".pkl")

    with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
        start_ndx = shard_id * num_per_shard
        end_ndx = min((shard_id + 1) * num_per_shard, len(filenames))
        for i in range(start_ndx, end_ndx):
            # sys.stdout.write('\r>> Converting image %d/%d shard %d' % (
            #     i + 1, len(filenames), shard_id))
            # sys.stdout.flush()

            # Read the filename:
            md5_seri.add_new_record(filenames[i])
            image_data = tf.gfile.FastGFile(filenames[i], 'rb').read()
            height, width = image_reader.read_image_dims(sess, image_data)

            class_name = os.path.basename(os.path.dirname(filenames[i]))
            class_id = class_names_to_ids[class_name]

            example = dataset_utils.image_to_tfexample(
                image_data, b'jpg', height, width, class_id)
            tfrecord_writer.write(example.SerializeToString())
    md5_seri.dump()
    g_process.add_finish(output_filename)
    sys.stdout.write('\r>> Converting image %d/%d of files: %d' % (
        len(g_process.finished_file), g_process.all_step, len(filenames)))
    sys.stdout.flush()


def thread_pool(enter_func, param_list):
    pool = threadpool.ThreadPool(cpu_count())
    requests = threadpool.makeRequests(enter_func, param_list)
    [pool.putRequest(req) for req in requests]
    pool.wait()


def _convert_dataset(split_name, filenames, class_names_to_ids, dataset_dir, num_shards):
    """Converts the given filenames to a TFRecord dataset.

    Args:
      split_name: The name of the dataset, either 'train' or 'validation'.
      filenames: A list of absolute paths to png or jpg images.
      class_names_to_ids: A dictionary from class names (strings) to ids
        (integers).
      dataset_dir: The directory where the converted datasets are stored.
    """
    assert split_name in ['train', 'validation']

    num_per_shard = int(math.ceil(len(filenames) / float(num_shards)))

    with tf.Graph().as_default():
        image_reader = ImageReader()

        with tf.Session('') as sess:
            param_list = []
            for shard_id in range(num_shards):
                # create_tf(split_name, filenames, class_names_to_ids, dataset_dir,
                #           num_shards, shard_id, num_per_shard, image_reader, sess)
                param_list.append(
                    ((split_name, filenames, class_names_to_ids, dataset_dir,
                      num_shards, shard_id, num_per_shard, image_reader, sess), None))
            thread_pool(create_tf, param_list)

    sys.stdout.write('\n')
    sys.stdout.flush()


def run(dataset_dir):
    """Runs the download and conversion operation.

    Args:
      dataset_dir: The dataset directory where the dataset is stored.
    """
    if not tf.gfile.Exists(dataset_dir):
        tf.gfile.MakeDirs(dataset_dir)

    photo_filenames, class_names = _get_filenames_and_classes(dataset_dir)
    if len(photo_filenames) == 0:
        print(" no files detected")
        exit(-1)
    zz = zip(class_names, range(len(class_names)))
    class_names_to_ids = dict(zip(class_names, range(len(class_names))))

    random.seed(_RANDOM_SEED)
    random.shuffle(photo_filenames)

    labels_to_class_names = dict(zip(range(len(class_names)), class_names))
    if check_label_same(OUTPUT_PATH + "/labels.txt", labels_to_class_names) is False:
        print("labels not found or changed, removing all ")
        if os.path.exists(OUTPUT_PATH) is True:
            shutil.rmtree(OUTPUT_PATH)
    if os.path.exists(OUTPUT_PATH) is False:
        os.makedirs(OUTPUT_PATH)
    if PREFIX == "train":
        dataset_utils.write_label_file(labels_to_class_names, OUTPUT_PATH)
    else:
        dataset_utils.write_label_file(labels_to_class_names, OUTPUT_PATH, PREFIX + "_labels.txt")

    print(' LOADING AT: %.5f second' % (time.time() - start_time))
    wipe_out_already_finished_task(photo_filenames, get_pkl_list(OUTPUT_PATH))
    check_pair_of_pkl_tf(OUTPUT_PATH)
    print(' CHECKING AT: %.5f second' % (time.time() - start_time))

    print("all files:%d classes: %d " % (len(photo_filenames), len(class_names)))
    num_shards = math.ceil(len(photo_filenames) / 1000)
    g_process.set_all_step(num_shards)
    if len(photo_filenames) == 0:
        return
    _convert_dataset(PREFIX, photo_filenames, class_names_to_ids,
                     dataset_dir, num_shards)


if __name__ == "__main__":
    start_time = time.time()
    VERSION = int(start_time)
    g_process = GlobalProcess()
    PREFIX = "train"
    # PREFIX = "validation"
    run("/home/leo/Downloads/train_datas/result_1/")
    # run("/home/leo/Downloads/train_datas/test_set_speed")
    # run("/home/leo/Downloads/train_datas/smallDataSetTest5_9/train_set")
    # run("/home/leo/Downloads/train_datas/smallDataSetTest5_9/validation_set")
    # run("../grabcut0410")
    print(' ALL COST: %.5f second' % (time.time() - start_time))
